#include "x25519.h"

#include <string.h> // memcpy

typedef unsigned __int128 uint128_t;

/*
h = f + g
Can overlap h with f or g.
*/

static inline void
fe25519_add(fe25519 h, const fe25519 f, const fe25519 g)
{
	uint64_t h0 = f[0] + g[0];
	uint64_t h1 = f[1] + g[1];
	uint64_t h2 = f[2] + g[2];
	uint64_t h3 = f[3] + g[3];
	uint64_t h4 = f[4] + g[4];

	h[0] = h0;
	h[1] = h1;
	h[2] = h2;
	h[3] = h3;
	h[4] = h4;
}

/*
h = f - g
*/

static void
fe25519_sub(fe25519 h, const fe25519 f, const fe25519 g)
{
	const uint64_t mask = 0x7ffffffffffffULL;
	uint64_t h0, h1, h2, h3, h4;

	h0 = g[0];
	h1 = g[1];
	h2 = g[2];
	h3 = g[3];
	h4 = g[4];

	h1 += h0 >> 51;
	h0 &= mask;
	h2 += h1 >> 51;
	h1 &= mask;
	h3 += h2 >> 51;
	h2 &= mask;
	h4 += h3 >> 51;
	h3 &= mask;
	h0 += 19ULL * (h4 >> 51);
	h4 &= mask;

	h0 = (f[0] + 0xfffffffffffdaULL) - h0;
	h1 = (f[1] + 0xffffffffffffeULL) - h1;
	h2 = (f[2] + 0xffffffffffffeULL) - h2;
	h3 = (f[3] + 0xffffffffffffeULL) - h3;
	h4 = (f[4] + 0xffffffffffffeULL) - h4;

	h[0] = h0;
	h[1] = h1;
	h[2] = h2;
	h[3] = h3;
	h[4] = h4;
}

/*
h = f * g
Can overlap h with f or g.
*/

extern "C" void x25519_fe51_mul(fe25519 h, const fe25519 f, const fe25519 g);
#define fe25519_mul x25519_fe51_mul

/*
static void
fe25519_mul(fe25519 h, const fe25519 f, const fe25519 g)
{
	const uint64_t mask = 0x7ffffffffffffULL;
	uint128_t r0, r1, r2, r3, r4, carry;
	uint64_t  f0, f1, f2, f3, f4;
	uint64_t  f1_19, f2_19, f3_19, f4_19;
	uint64_t  g0, g1, g2, g3, g4;
	uint64_t  r00, r01, r02, r03, r04;

	f0 = f[0];
	f1 = f[1];
	f2 = f[2];
	f3 = f[3];
	f4 = f[4];

	g0 = g[0];
	g1 = g[1];
	g2 = g[2];
	g3 = g[3];
	g4 = g[4];

	f1_19 = 19ULL * f1;
	f2_19 = 19ULL * f2;
	f3_19 = 19ULL * f3;
	f4_19 = 19ULL * f4;

	r0 = ((uint128_t)f0) * ((uint128_t)g0);
	r0 += ((uint128_t)f1_19) * ((uint128_t)g4);
	r0 += ((uint128_t)f2_19) * ((uint128_t)g3);
	r0 += ((uint128_t)f3_19) * ((uint128_t)g2);
	r0 += ((uint128_t)f4_19) * ((uint128_t)g1);

	r1 = ((uint128_t)f0) * ((uint128_t)g1);
	r1 += ((uint128_t)f1) * ((uint128_t)g0);
	r1 += ((uint128_t)f2_19) * ((uint128_t)g4);
	r1 += ((uint128_t)f3_19) * ((uint128_t)g3);
	r1 += ((uint128_t)f4_19) * ((uint128_t)g2);

	r2 = ((uint128_t)f0) * ((uint128_t)g2);
	r2 += ((uint128_t)f1) * ((uint128_t)g1);
	r2 += ((uint128_t)f2) * ((uint128_t)g0);
	r2 += ((uint128_t)f3_19) * ((uint128_t)g4);
	r2 += ((uint128_t)f4_19) * ((uint128_t)g3);

	r3 = ((uint128_t)f0) * ((uint128_t)g3);
	r3 += ((uint128_t)f1) * ((uint128_t)g2);
	r3 += ((uint128_t)f2) * ((uint128_t)g1);
	r3 += ((uint128_t)f3) * ((uint128_t)g0);
	r3 += ((uint128_t)f4_19) * ((uint128_t)g4);

	r4 = ((uint128_t)f0) * ((uint128_t)g4);
	r4 += ((uint128_t)f1) * ((uint128_t)g3);
	r4 += ((uint128_t)f2) * ((uint128_t)g2);
	r4 += ((uint128_t)f3) * ((uint128_t)g1);
	r4 += ((uint128_t)f4) * ((uint128_t)g0);

	r00 = ((uint64_t)r0) & mask;
	carry = r0 >> 51;
	r1 += carry;
	r01 = ((uint64_t)r1) & mask;
	carry = r1 >> 51;
	r2 += carry;
	r02 = ((uint64_t)r2) & mask;
	carry = r2 >> 51;
	r3 += carry;
	r03 = ((uint64_t)r3) & mask;
	carry = r3 >> 51;
	r4 += carry;
	r04 = ((uint64_t)r4) & mask;
	carry = r4 >> 51;
	r00 += 19ULL * (uint64_t)carry;
	carry = r00 >> 51;
	r00 &= mask;
	r01 += (uint64_t)carry;
	carry = r01 >> 51;
	r01 &= mask;
	r02 += (uint64_t)carry;

	h[0] = r00;
	h[1] = r01;
	h[2] = r02;
	h[3] = r03;
	h[4] = r04;
}
*/

/*
h = f * f
Can overlap h with f.
*/

extern "C" void x25519_fe51_sqr(fe25519 h, const fe25519 f);
#define fe25519_sq x25519_fe51_sqr

/*
static void
fe25519_sq(fe25519 h, const fe25519 f)
{
	const uint64_t mask = 0x7ffffffffffffULL;
	uint128_t r0, r1, r2, r3, r4, carry;
	uint64_t  f0, f1, f2, f3, f4;
	uint64_t  f0_2, f1_2, f1_38, f2_38, f3_38, f3_19, f4_19;
	uint64_t  r00, r01, r02, r03, r04;

	f0 = f[0];
	f1 = f[1];
	f2 = f[2];
	f3 = f[3];
	f4 = f[4];

	f0_2 = f0 << 1;
	f1_2 = f1 << 1;

	f1_38 = 38ULL * f1;
	f2_38 = 38ULL * f2;
	f3_38 = 38ULL * f3;

	f3_19 = 19ULL * f3;
	f4_19 = 19ULL * f4;

	r0 = ((uint128_t)f0) * ((uint128_t)f0);
	r0 += ((uint128_t)f1_38) * ((uint128_t)f4);
	r0 += ((uint128_t)f2_38) * ((uint128_t)f3);

	r1 = ((uint128_t)f0_2) * ((uint128_t)f1);
	r1 += ((uint128_t)f2_38) * ((uint128_t)f4);
	r1 += ((uint128_t)f3_19) * ((uint128_t)f3);

	r2 = ((uint128_t)f0_2) * ((uint128_t)f2);
	r2 += ((uint128_t)f1) * ((uint128_t)f1);
	r2 += ((uint128_t)f3_38) * ((uint128_t)f4);

	r3 = ((uint128_t)f0_2) * ((uint128_t)f3);
	r3 += ((uint128_t)f1_2) * ((uint128_t)f2);
	r3 += ((uint128_t)f4_19) * ((uint128_t)f4);

	r4 = ((uint128_t)f0_2) * ((uint128_t)f4);
	r4 += ((uint128_t)f1_2) * ((uint128_t)f3);
	r4 += ((uint128_t)f2) * ((uint128_t)f2);

	r00 = ((uint64_t)r0) & mask;
	carry = r0 >> 51;
	r1 += carry;
	r01 = ((uint64_t)r1) & mask;
	carry = r1 >> 51;
	r2 += carry;
	r02 = ((uint64_t)r2) & mask;
	carry = r2 >> 51;
	r3 += carry;
	r03 = ((uint64_t)r3) & mask;
	carry = r3 >> 51;
	r4 += carry;
	r04 = ((uint64_t)r4) & mask;
	carry = r4 >> 51;
	r00 += 19ULL * (uint64_t)carry;
	carry = r00 >> 51;
	r00 &= mask;
	r01 += (uint64_t)carry;
	carry = r01 >> 51;
	r01 &= mask;
	r02 += (uint64_t)carry;

	h[0] = r00;
	h[1] = r01;
	h[2] = r02;
	h[3] = r03;
	h[4] = r04;
}
*/

static void
fe25519_invert(fe25519 out, const fe25519 z)
{
	fe25519 t0;
	fe25519 t1;
	fe25519 t2;
	fe25519 t3;
	int     i;

	fe25519_sq(t0, z);
	fe25519_sq(t1, t0);
	fe25519_sq(t1, t1);
	fe25519_mul(t1, z, t1);
	fe25519_mul(t0, t0, t1);
	fe25519_sq(t2, t0);
	fe25519_mul(t1, t1, t2);
	fe25519_sq(t2, t1);
	for (i = 1; i < 5; ++i) {
		fe25519_sq(t2, t2);
	}
	fe25519_mul(t1, t2, t1);
	fe25519_sq(t2, t1);
	for (i = 1; i < 10; ++i) {
		fe25519_sq(t2, t2);
	}
	fe25519_mul(t2, t2, t1);
	fe25519_sq(t3, t2);
	for (i = 1; i < 20; ++i) {
		fe25519_sq(t3, t3);
	}
	fe25519_mul(t2, t3, t2);
	for (i = 1; i < 11; ++i) {
		fe25519_sq(t2, t2);
	}
	fe25519_mul(t1, t2, t1);
	fe25519_sq(t2, t1);
	for (i = 1; i < 50; ++i) {
		fe25519_sq(t2, t2);
	}
	fe25519_mul(t2, t2, t1);
	fe25519_sq(t3, t2);
	for (i = 1; i < 100; ++i) {
		fe25519_sq(t3, t3);
	}
	fe25519_mul(t2, t3, t2);
	for (i = 1; i < 51; ++i) {
		fe25519_sq(t2, t2);
	}
	fe25519_mul(t1, t2, t1);
	for (i = 1; i < 6; ++i) {
		fe25519_sq(t1, t1);
	}
	fe25519_mul(out, t1, t0);
}

static void
fe25519_reduce(fe25519 h, const fe25519 f)
{
	const uint64_t mask = 0x7ffffffffffffULL;
	uint128_t t[5];

	t[0] = f[0];
	t[1] = f[1];
	t[2] = f[2];
	t[3] = f[3];
	t[4] = f[4];

	t[1] += t[0] >> 51;
	t[0] &= mask;
	t[2] += t[1] >> 51;
	t[1] &= mask;
	t[3] += t[2] >> 51;
	t[2] &= mask;
	t[4] += t[3] >> 51;
	t[3] &= mask;
	t[0] += 19 * (t[4] >> 51);
	t[4] &= mask;

	t[1] += t[0] >> 51;
	t[0] &= mask;
	t[2] += t[1] >> 51;
	t[1] &= mask;
	t[3] += t[2] >> 51;
	t[2] &= mask;
	t[4] += t[3] >> 51;
	t[3] &= mask;
	t[0] += 19 * (t[4] >> 51);
	t[4] &= mask;

	/* now t is between 0 and 2^255-1, properly carried. */
	/* case 1: between 0 and 2^255-20. case 2: between 2^255-19 and 2^255-1. */

	t[0] += 19ULL;

	t[1] += t[0] >> 51;
	t[0] &= mask;
	t[2] += t[1] >> 51;
	t[1] &= mask;
	t[3] += t[2] >> 51;
	t[2] &= mask;
	t[4] += t[3] >> 51;
	t[3] &= mask;
	t[0] += 19ULL * (t[4] >> 51);
	t[4] &= mask;

	/* now between 19 and 2^255-1 in both cases, and offset by 19. */

	t[0] += 0x8000000000000 - 19ULL;
	t[1] += 0x8000000000000 - 1ULL;
	t[2] += 0x8000000000000 - 1ULL;
	t[3] += 0x8000000000000 - 1ULL;
	t[4] += 0x8000000000000 - 1ULL;

	/* now between 2^255 and 2^256-20, and offset by 2^255. */

	t[1] += t[0] >> 51;
	t[0] &= mask;
	t[2] += t[1] >> 51;
	t[1] &= mask;
	t[3] += t[2] >> 51;
	t[2] &= mask;
	t[4] += t[3] >> 51;
	t[3] &= mask;
	t[4] &= mask;

	h[0] = t[0];
	h[1] = t[1];
	h[2] = t[2];
	h[3] = t[3];
	h[4] = t[4];
}

static void
fe25519_tobytes(unsigned char *s, const fe25519 h)
{
	fe25519  t;
	uint64_t t0, t1, t2, t3;

	fe25519_reduce(t, h);
	t0 = t[0] | (t[1] << 51);
	t1 = (t[1] >> 13) | (t[2] << 38);
	t2 = (t[2] >> 26) | (t[3] << 25);
	t3 = (t[3] >> 39) | (t[4] << 12);
	memcpy(s + 0, &t0, sizeof t0);
	memcpy(s + 8, &t1, sizeof t1);
	memcpy(s + 16, &t2, sizeof t2);
	memcpy(s + 24, &t3, sizeof t3);
}

static void ge25519_add(ge25519* r, const ge25519* p, const ge25519* q)
{
	static const fe25519 d2 = {
		0x69b9426b2f159, 0x35050762add7a,
		0x3cf44c0038052, 0x6738cc7407977, 0x2406d9dc56dff };

	fe25519 a, b, c, d, t, e, f, g, h;

	fe25519_sub(a, p->Y, p->X);
	fe25519_sub(t, q->Y, q->X);
	fe25519_mul(a, a, t);
	fe25519_add(b, p->X, p->Y);
	fe25519_add(t, q->X, q->Y);
	fe25519_mul(b, b, t);
	fe25519_mul(c, p->T, q->T);
	fe25519_mul(c, c, d2);
	fe25519_mul(d, p->Z, q->Z);
	fe25519_add(d, d, d);
	fe25519_sub(e, b, a);
	fe25519_sub(f, d, c);
	fe25519_add(g, d, c);
	fe25519_add(h, b, a);

	fe25519_mul(r->X, e, f);
	fe25519_mul(r->Y, h, g);
	fe25519_mul(r->Z, g, f);
	fe25519_mul(r->T, e, h);
}


void x25519::ge25519_scalarmult_base(ge25519* h)
{
	if (state.size() == 0)
		*h = basePowers[254];
	else
		*h = state.back();

	for (int i = 253 - state.size(); i >= 3; i--)
	{
		if (keyBits[i])
			ge25519_add(h, h, &basePowers[i]);
		state.push_back(*h);
	}
}

static void edwards_to_montgomery(fe25519 montgomeryX,
	const fe25519 edwardsY, const fe25519 edwardsZ)
{
	fe25519 tempX;
	fe25519 tempZ;

	fe25519_add(tempX, edwardsZ, edwardsY);
	fe25519_sub(tempZ, edwardsZ, edwardsY);
	fe25519_invert(tempZ, tempZ);
	fe25519_mul(montgomeryX, tempX, tempZ);
}

x25519::x25519()
{
	ge25519 q = {
		{ 0x62d608f25d51a, 0x412a4b4f6592a,
		  0x75b7171a4b31d, 0x1ff60527118fe, 0x216936d3cd6e5 },
		{ 0x6666666666658, 0x4cccccccccccc,
		  0x1999999999999, 0x3333333333333, 0x6666666666666 },
		{ 1, 0, 0, 0, 0 },
		{ 0x68ab3a5b7dda3, 0x00eea2a5eadbb, 0x2af8df483c27e,
		  0x332b375274732, 0x67875f0fd78b7 }
	};

	for (int i = 0; i < 255; i++)
	{
		basePowers[i] = q;
		ge25519_add(&q, &q, &q);
	}
}

void x25519::init_private_key(const uint8_t privateKey[KEYSIZE])
{
	for (int i = 0; i < 256; i++)
		keyBits[i] = (privateKey[i / 8] >> (i & 7)) & 1;
	state.clear();
}

void x25519::mutate_private_key()
{
	for (int i = 3; ; i++)
	{
		keyBits[i]--;
		state.pop_back();
		if (keyBits[i] == 0)
			break;
		else
			keyBits[i] = 1;
	}
}

void x25519::get_private_key(uint8_t privateKey[KEYSIZE])
{
	for (int i = 0; i < 32; i++)
		privateKey[i] = 0;
	for (int i = 0; i < 256; i++)
		privateKey[i / 8] |= keyBits[i] << (i & 7);
}

void x25519::get_public_key(uint8_t publicKey[KEYSIZE])
{
	ge25519 a;
	fe25519 pk;
	ge25519_scalarmult_base(&a);
	edwards_to_montgomery(pk, a.Y, a.Z);
	fe25519_tobytes(publicKey, pk);
}